package org.example.algorithm.trie;

import java.util.ArrayList;
import java.util.List;
import java.util.Stack;

/**
 * 使用trie树+栈
 */
public class WordBreakTrie {
    private WordBreakTrie[] child;
    private boolean isEnd;

    public WordBreakTrie() {
        child = new WordBreakTrie[26];
        isEnd = false;
    }

    public void insert(String word) {
        WordBreakTrie trie = this;
        for (int i=0;i<word.length();i++) {
            int index = word.charAt(i) - 'a';
            if (trie.child[index] == null) {
                trie.child[index] = new WordBreakTrie();
            }
            trie = trie.child[index];
        }
        trie.isEnd = true;
    }

    public boolean wordBreak(String s, List<String> wordDict) {
        int len = s.length();
        if (len == 0) {
            return true;
        }
        for (String word: wordDict) {
            insert(word);
        }
        Stack<Integer> stack = new Stack<>();
        stack.push(0);
        boolean[] search = new boolean[len+1];
        search[0] = true;
        while (!stack.isEmpty()) {
            WordBreakTrie trie = this;
            int index = stack.pop();
            while (index<len && trie.child[s.charAt(index) - 'a'] != null) {
                int trieIndex = s.charAt(index) - 'a';
                index++;
                trie = trie.child[trieIndex];
                if (trie.isEnd && !search[index]) {
                    stack.push(index);
                    search[index] = true;
                }
            }
            if (index == len && trie.isEnd) {
                return true;
            }
        }
        return false;
    }

    public static void main(String[] args) {
        String s = "aaab";
        List<String> wordDict = new ArrayList<>();
        wordDict.add("a");
        wordDict.add("aa");
        wordDict.add("aaa");
        WordBreakTrie wordBreakTrie = new WordBreakTrie();
        boolean res = wordBreakTrie.wordBreak(s, wordDict);
        System.out.println(res);
    }
}
